package org.example.myleet.p726;

import java.util.*;

public class Solution {

    private StringBuilder atomName;

    private int count;

    private int readIdx;

    private Deque<Map<String, Integer>> stack;

    public String countOfAtoms(String formula) {
        Map<String, Integer> result = new HashMap<>();
        stack = new LinkedList<>();
        stack.push(result);
        readIdx = 0;
        atomName = new StringBuilder();
        count = 0;
        while (readIdx < formula.length()) {
            char c = formula.charAt(readIdx++);
            if (c == '(') {
                if (atomName.length() != 0) {
                    //找到一个原子组成
                    processAtom();
                }
                //栈新增下一层map进行处理
                Map<String, Integer> subMap = new HashMap<>();
                stack.push(subMap);
            } else if (c == ')') {
                if (atomName.length() != 0) {
                    //找到一个原子组成
                    processAtom();
                }
                //栈弹出一层map并合并到当前层的map中
                //1、先找到右括号后的数字
                int complexCount = 0;
                while (readIdx < formula.length()) {
                    c = formula.charAt(readIdx++);
                    if (Character.isDigit(c)) {
                        complexCount = complexCount * 10 + (c - '0');
                    } else {
                        --readIdx;
                        break;
                    }
                }
                complexCount = complexCount > 0 ? complexCount : 1;
                //2、进行map合并
                Map<String, Integer> complexAtomMap = stack.pop();
                Map<String, Integer> atomMap = stack.peek();
                for (Map.Entry<String, Integer> entry : complexAtomMap.entrySet()) {
                    Integer atomCount = atomMap.getOrDefault(entry.getKey(), 0);
                    atomCount += entry.getValue() * complexCount;
                    atomMap.put(entry.getKey(), atomCount);
                }
            } else if (Character.isDigit(c)) {
                count = count * 10 + (c - '0');
            } else if ('A' <= c && c <= 'Z') {
                if (atomName.length() != 0) {
                    //找到一个原子组成
                    processAtom();
                    atomName.append(c);
                } else {
                    //新建一个原子名称
                    atomName.append(c);
                }
            } else if ('a' <= c && c <= 'z') {
                atomName.append(c);
            }
        }
        if (atomName.length() != 0) {
            //找到一个原子组成
            processAtom();
        }
        Map<String, Integer> atomMap = stack.pop();
        //通过堆排序输出原子名称字典序排序
        PriorityQueue<Map.Entry<String, Integer>> heap = new PriorityQueue<>(Map.Entry.comparingByKey());
        for (Map.Entry<String, Integer> entry : atomMap.entrySet()) {
            heap.offer(entry);
        }
        StringBuilder simple = new StringBuilder();
        while (!heap.isEmpty()) {
            Map.Entry<String, Integer> entry = heap.poll();
            simple.append(entry.getKey());
            if (entry.getValue() > 1) {
                simple.append(entry.getValue());
            }
        }
        return simple.toString();
    }

    /**
     * 找到一个原子组成，处理成原子数目
     */
    private void processAtom() {
        count = count > 0 ? count : 1;
        String name = atomName.toString();
        Map<String, Integer> atomMap = stack.peek();
        Integer atomCount = atomMap.getOrDefault(name, 0);
        atomCount += count;
        atomMap.put(name, atomCount);
        atomName.delete(0, atomName.length());
        count = 0;
    }
}
